import random
from enum import Enum
import time

class Action(Enum):
    UP = 'UP'
    DOWN = 'DOWN'
    LEFT = 'LEFT'
    RIGHT = 'RIGHT'
    GRAB = 'GRAB'
    SHOOT = 'SHOOT'
    CLIMB = 'CLIMB'

class WumpusWorld:
    def __init__(self, size=4, random_world=False):
        self.size = size
        self.agent_position = (0, 0)
        self.agent_direction = 'RIGHT'
        self.agent_alive = True
        self.has_gold = False
        self.agent_has_arrow = True
        self.wumpus_alive = True
        
        # 初始化世界
        self.grid = [[{'pit': False, 'wumpus': False, 'gold': False, 
                      'breeze': False, 'stench': False, 'glitter': False} 
                     for _ in range(size)] for _ in range(size)]
        
        if random_world:
            self._create_random_world()
        else:
            # 创建标准的Wumpus世界
            self._create_standard_world()
        
    def _create_standard_world(self):
        # 放置Wumpus怪物
        self.grid[0][3]['wumpus'] = True
        
        # 放置陷阱
        self.grid[2][0]['pit'] = True
        self.grid[1][2]['pit'] = True
        self.grid[3][3]['pit'] = True
        
        # 放置黄金
        self.grid[1][1]['gold'] = True
        self.grid[1][1]['glitter'] = True
        
        # 更新感知
        self._update_percepts()
    
    def _create_random_world(self):
        # 确保起点安全
        safe_positions = [(0, 0)]
        
        # 随机放置Wumpus（不在起点）
        while True:
            x, y = random.randint(0, self.size-1), random.randint(0, self.size-1)
            if (x, y) not in safe_positions:
                self.grid[x][y]['wumpus'] = True
                break
        
        # 随机放置陷阱（不在起点，概率0.2）
        for x in range(self.size):
            for y in range(self.size):
                if (x, y) not in safe_positions and random.random() < 0.2:
                    self.grid[x][y]['pit'] = True
        
        # 随机放置黄金（不在起点）
        while True:
            x, y = random.randint(0, self.size-1), random.randint(0, self.size-1)
            if (x, y) not in safe_positions and not self.grid[x][y]['wumpus'] and not self.grid[x][y]['pit']:
                self.grid[x][y]['gold'] = True
                self.grid[x][y]['glitter'] = True
                break
        
        # 更新感知
        self._update_percepts()
    
    def _update_percepts(self):
        # 更新臭气和微风
        for x in range(self.size):
            for y in range(self.size):
                if self.grid[x][y]['wumpus']:
                    # 添加臭气到相邻格子
                    for nx, ny in self._get_neighbors(x, y):
                        self.grid[nx][ny]['stench'] = True
                
                if self.grid[x][y]['pit']:
                    # 添加微风到相邻格子
                    for nx, ny in self._get_neighbors(x, y):
                        self.grid[nx][ny]['breeze'] = True
    
    def _get_neighbors(self, x, y):
        neighbors = []
        for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
            nx, ny = x + dx, y + dy
            if 0 <= nx < self.size and 0 <= ny < self.size:
                neighbors.append((nx, ny))
        return neighbors
    
    def get_percept(self):
        x, y = self.agent_position
        if not self.agent_alive:
            return None
        
        return {
            'stench': self.grid[x][y]['stench'],
            'breeze': self.grid[x][y]['breeze'],
            'glitter': self.grid[x][y]['glitter'],
            'bump': False,  # 初始没有撞击感知
            'scream': False  # 初始没有尖叫感知
        }
    
    def execute_action(self, action):
        if not self.agent_alive:
            return False
        
        x, y = self.agent_position
        
        if action == Action.UP:
            if y < self.size - 1:
                self.agent_position = (x, y + 1)
            else:
                return {'bump': True}
        
        elif action == Action.DOWN:
            if y > 0:
                self.agent_position = (x, y - 1)
            else:
                return {'bump': True}
        
        elif action == Action.LEFT:
            if x > 0:
                self.agent_position = (x - 1, y)
            else:
                return {'bump': True}
        
        elif action == Action.RIGHT:
            if x < self.size - 1:
                self.agent_position = (x + 1, y)
            else:
                return {'bump': True}
        
        elif action == Action.GRAB:
            nx, ny = self.agent_position
            if self.grid[nx][ny]['gold']:
                self.has_gold = True
                self.grid[nx][ny]['gold'] = False
                self.grid[nx][ny]['glitter'] = False
        
        elif action == Action.SHOOT:
            if self.agent_has_arrow:
                self.agent_has_arrow = False
                # 检查是否击中Wumpus
                nx, ny = self.agent_position
                dx, dy = {'RIGHT': (1, 0), 'LEFT': (-1, 0), 'UP': (0, 1), 'DOWN': (0, -1)}[self.agent_direction]
                
                while 0 <= nx < self.size and 0 <= ny < self.size:
                    nx += dx
                    ny += dy
                    if 0 <= nx < self.size and 0 <= ny < self.size and self.grid[nx][ny]['wumpus']:
                        self.wumpus_alive = False
                        self.grid[nx][ny]['wumpus'] = False
                        # 更新臭气
                        for rx, ry in self._get_neighbors(nx, ny):
                            has_other_wumpus = False
                            for wx, wy in self._get_neighbors(rx, ry):
                                if (wx, wy) != (nx, ny) and self.grid[wx][wy]['wumpus']:
                                    has_other_wumpus = True
                                    break
                            if not has_other_wumpus:
                                self.grid[rx][ry]['stench'] = False
                        return {'scream': True}
        
        elif action == Action.CLIMB:
            if self.agent_position == (0, 0) and self.has_gold:
                return {'success': True}
        
        # 检查是否踩到陷阱或Wumpus
        nx, ny = self.agent_position
        if self.grid[nx][ny]['pit'] or (self.grid[nx][ny]['wumpus'] and self.wumpus_alive):
            self.agent_alive = False
            return {'dead': True}
        
        return self.get_percept()

class Agent:
    def __init__(self, world_size):
        self.world_size = world_size
        self.position = (0, 0)
        self.direction = 'RIGHT'
        self.has_gold = False
        self.has_arrow = True
        
        # 知识库
        self.KB = [[{'visited': False, 'safe': False, 'wumpus': False, 'pit': False, 
                     'wumpus_possible': False, 'pit_possible': False} 
                   for _ in range(world_size)] for _ in range(world_size)]
        
        # 标记起点为安全和已访问
        self.KB[0][0]['safe'] = True
        self.KB[0][0]['visited'] = True
        
        # 联机搜索状态
        self.path_home = []  # 回家的路径
        self.frontier = []   # 待探索的边界
        self.add_frontier_cells(0, 0)
    
    def add_frontier_cells(self, x, y):
        for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
            nx, ny = x + dx, y + dy
            if 0 <= nx < self.world_size and 0 <= ny < self.world_size and not self.KB[nx][ny]['visited']:
                if (nx, ny) not in self.frontier:
                    self.frontier.append((nx, ny))
    
    def update_knowledge(self, percept):
        x, y = self.position
        
        # 标记当前位置为安全和已访问
        self.KB[x][y]['safe'] = True
        self.KB[x][y]['visited'] = True
        self.KB[x][y]['wumpus_possible'] = False
        self.KB[x][y]['pit_possible'] = False
        
        # 根据感知更新知识
        if percept.get('stench'):
            # 可能有Wumpus在周围
            for nx, ny in self.get_neighbors(x, y):
                if not self.KB[nx][ny]['visited']:
                    self.KB[nx][ny]['wumpus_possible'] = True
        
        if percept.get('breeze'):
            # 可能有陷阱在周围
            for nx, ny in self.get_neighbors(x, y):
                if not self.KB[nx][ny]['visited']:
                    self.KB[nx][ny]['pit_possible'] = True
        
        # 如果既没有臭气也没有微风，周围格子是安全的
        if not percept.get('stench') and not percept.get('breeze'):
            for nx, ny in self.get_neighbors(x, y):
                self.KB[nx][ny]['safe'] = True
                self.KB[nx][ny]['wumpus_possible'] = False
                self.KB[nx][ny]['pit_possible'] = False
        
        # 更新边界
        self.add_frontier_cells(x, y)
    
    def get_neighbors(self, x, y):
        neighbors = []
        for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
            nx, ny = x + dx, y + dy
            if 0 <= nx < self.world_size and 0 <= ny < self.world_size:
                neighbors.append((nx, ny))
        return neighbors
    
    def choose_action(self, percept):
        x, y = self.position
        
        # 更新知识库
        self.update_knowledge(percept)
        
        # 如果当前位置有黄金，抓取它
        if percept.get('glitter'):
            self.has_gold = True
            return Action.GRAB
        
        # 如果有黄金并且在起点，离开洞穴
        if self.has_gold and self.position == (0, 0):
            return Action.CLIMB
        
        # 如果有黄金，尝试回家
        if self.has_gold:
            if not self.path_home:
                # 计算回家的路径
                self.path_home = self.find_path_to((0, 0))
            
            if self.path_home:
                next_pos = self.path_home.pop(0)
                dx, dy = next_pos[0] - x, next_pos[1] - y
                if dx == 1:
                    self.direction = 'RIGHT'
                    return Action.RIGHT
                elif dx == -1:
                    self.direction = 'LEFT'
                    return Action.LEFT
                elif dy == 1:
                    self.direction = 'UP'
                    return Action.UP
                else:
                    self.direction = 'DOWN'
                    return Action.DOWN
        
        # 探索安全的未访问格子
        safe_unvisited = []
        for nx in range(self.world_size):
            for ny in range(self.world_size):
                if self.KB[nx][ny]['safe'] and not self.KB[nx][ny]['visited']:
                    safe_unvisited.append((nx, ny))
        
        if safe_unvisited:
            # 找到最近的安全未访问格子
            path = self.find_path_to(min(safe_unvisited, key=lambda pos: abs(pos[0]-x) + abs(pos[1]-y)))
            if path:
                next_pos = path[0]
                dx, dy = next_pos[0] - x, next_pos[1] - y
                if dx == 1:
                    self.direction = 'RIGHT'
                    return Action.RIGHT
                elif dx == -1:
                    self.direction = 'LEFT'
                    return Action.LEFT
                elif dy == 1:
                    self.direction = 'UP'
                    return Action.UP
                else:
                    self.direction = 'DOWN'
                    return Action.DOWN
        
        # 如果没有安全的未访问格子，冒险尝试可能安全的格子
        if self.frontier:
            # 按照启发式计算风险最低的边界格子
            best_frontier = None
            lowest_risk = float('inf')
            
            for fx, fy in self.frontier:
                risk = 0
                if self.KB[fx][fy]['wumpus_possible']:
                    risk += 10
                if self.KB[fx][fy]['pit_possible']:
                    risk += 5
                
                # 如果我们还有箭并且怀疑有Wumpus，风险较低
                if self.KB[fx][fy]['wumpus_possible'] and self.has_arrow:
                    risk -= 8
                
                # 距离也算入风险
                risk += abs(fx-x) + abs(fy-y)
                
                if risk < lowest_risk:
                    lowest_risk = risk
                    best_frontier = (fx, fy)
            
            if best_frontier:
                # 如果目标格子可能有Wumpus且我们有箭，先尝试射击
                if self.KB[best_frontier[0]][best_frontier[1]]['wumpus_possible'] and self.has_arrow:
                    # 调整方向准备射击
                    tx, ty = best_frontier
                    if tx > x and self.direction != 'RIGHT':
                        self.direction = 'RIGHT'
                    elif tx < x and self.direction != 'LEFT':
                        self.direction = 'LEFT'
                    elif ty > y and self.direction != 'UP':
                        self.direction = 'UP'
                    elif ty < y and self.direction != 'DOWN':
                        self.direction = 'DOWN'
                    else:
                        # 方向已对准，射击
                        self.has_arrow = False
                        return Action.SHOOT
                else:
                    # 移动到风险最低的边界格子
                    path = self.find_path_to(best_frontier)
                    if path:
                        next_pos = path[0]
                        dx, dy = next_pos[0] - x, next_pos[1] - y
                        if dx == 1:
                            self.direction = 'RIGHT'
                            return Action.RIGHT
                        elif dx == -1:
                            self.direction = 'LEFT'
                            return Action.LEFT
                        elif dy == 1:
                            self.direction = 'UP'
                            return Action.UP
                        else:
                            self.direction = 'DOWN'
                            return Action.DOWN
        
        # 如果实在没有选择，回到起点
        path = self.find_path_to((0, 0))
        if path:
            next_pos = path[0]
            dx, dy = next_pos[0] - x, next_pos[1] - y
            if dx == 1:
                self.direction = 'RIGHT'
                return Action.RIGHT
            elif dx == -1:
                self.direction = 'LEFT'
                return Action.LEFT
            elif dy == 1:
                self.direction = 'UP'
                return Action.UP
            else:
                self.direction = 'DOWN'
                return Action.DOWN
        
        # 如果所有策略都失败，随机行动
        return random.choice([Action.UP, Action.DOWN, Action.LEFT, Action.RIGHT])
    
    def find_path_to(self, target):
        """使用A*算法寻找从当前位置到目标位置的安全路径"""
        start = self.position
        if start == target:
            return []
        
        open_list = [(start, [])]  # (位置, 路径)
        closed_set = set()
        
        while open_list:
            (x, y), path = open_list.pop(0)
            
            if (x, y) == target:
                return path
            
            if (x, y) in closed_set:
                continue
            
            closed_set.add((x, y))
            
            for nx, ny in self.get_neighbors(x, y):
                if (nx, ny) not in closed_set and self.KB[nx][ny]['safe']:
                    new_path = path + [(nx, ny)]
                    # 启发式函数：曼哈顿距离
                    h = abs(nx - target[0]) + abs(ny - target[1])
                    # 按f = g + h排序，其中g是路径长度
                    i = 0
                    while i < len(open_list) and len(new_path) + h > len(open_list[i][1]) + abs(open_list[i][0][0] - target[0]) + abs(open_list[i][0][1] - target[1]):
                        i += 1
                    open_list.insert(i, ((nx, ny), new_path))
        
        return []  # 没有找到路径
    
    def execute_action(self, action, percept_result):
        # 更新智能体状态
        if action == Action.UP:
            if not percept_result.get('bump'):
                self.position = (self.position[0], self.position[1] + 1)
        elif action == Action.DOWN:
            if not percept_result.get('bump'):
                self.position = (self.position[0], self.position[1] - 1)
        elif action == Action.LEFT:
            if not percept_result.get('bump'):
                self.position = (self.position[0] - 1, self.position[1])
        elif action == Action.RIGHT:
            if not percept_result.get('bump'):
                self.position = (self.position[0] + 1, self.position[1])

def print_world(world, agent=None):
    """打印当前世界状态"""
    size = world.size
    
    # 打印顶部边界
    print('+'+ '-'*size*4 + '+')
    
    for y in range(size-1, -1, -1):
        row = '|'
        for x in range(size):
            cell = ' '
            if agent and agent.position == (x, y):
                cell = 'A'
            elif world.grid[x][y]['wumpus'] and world.wumpus_alive:
                cell = 'W'
            elif world.grid[x][y]['pit']:
                cell = 'P'
            elif world.grid[x][y]['gold']:
                cell = 'G'
            
            row += ' ' + cell + ' |'
        print(row)
        print('+'+ '-'*size*4 + '+')
    
    # 打印智能体状态
    if agent:
        print(f"智能体位置: {agent.position}, 方向: {agent.direction}")
        print(f"拥有黄金: {agent.has_gold}, 拥有箭: {agent.has_arrow}")

def print_knowledge(agent):
    """打印智能体的知识库"""
    size = agent.world_size
    
    print("\n智能体知识库:")
    print('+'+ '-'*size*8 + '+')
    
    for y in range(size-1, -1, -1):
        row = '|'
        for x in range(size):
            cell = ''
            if agent.KB[x][y]['visited']:
                cell += 'V'
            if agent.KB[x][y]['safe']:
                cell += 'S'
            if agent.KB[x][y]['wumpus_possible']:
                cell += 'W?'
            if agent.KB[x][y]['pit_possible']:
                cell += 'P?'
            
            row += ' ' + cell.ljust(6) + ' |'
        print(row)
        print('+'+ '-'*size*8 + '+')

def run_simulation():
    # 创建Wumpus世界
    world = WumpusWorld(size=4)
    
    # 创建智能体
    agent = Agent(world_size=4)
    
    # 模拟执行
    max_steps = 100
    step = 0
    game_over = False
    
    print("初始状态:")
    print_world(world, agent)
    
    while step < max_steps and not game_over:
        print(f"\n步骤 {step + 1}:")
        
        # 获取当前感知
        percept = world.get_percept()
        print(f"感知: {percept}")
        
        # 智能体选择行动
        action = agent.choose_action(percept)
        print(f"行动: {action}")
        
        # 执行行动
        result = world.execute_action(action)
        agent.execute_action(action, result)
        
        # 打印当前状态
        print_world(world, agent)
        print_knowledge(agent)
        
        # 检查游戏是否结束
        if result.get('success'):
            print("\n成功! 智能体找到黄金并安全返回。")
            game_over = True
        elif result.get('dead'):
            print("\n失败! 智能体死亡。")
            game_over = True
        
        step += 1
        time.sleep(0.5)  # 延迟以便观察
    
    if step >= max_steps:
        print("\n达到最大步数，模拟结束。")

if __name__ == "__main__":
    run_simulation()